import pygame
import numpy as np
import tensorflow as tf
from me_package import draw_stroke_rect, amplitude_reduction

# 加载模型
model = tf.keras.models.load_model("path_to_save_model.h5")

# 初始化Pygame
pygame.init()
# 定义窗口大小
window_width = 20 * 28
window_height = 20 * 28
# 创建窗口对象
window = pygame.display.set_mode((window_width + 200, window_height))
pygame.display.set_caption("手写数字识别")
# 创建矩形
rect_size = [20, 20]
rect_outline = 1
color_white = [255, 255, 255]
color_black = [0, 0, 0]
color_grey = [128, 128, 128]
# 定义字体和字体大小位置
predicted_label = 0
probability = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
font = pygame.font.Font(None, 30)
font_num1 = pygame.font.Font(None, 60)
font_num2 = pygame.font.Font(None, 30)
# 主循环
pixel_size = (28, 28)
pixel = np.zeros(pixel_size)
test_size = (28, 28)
test = np.zeros(test_size)

make = 0
running = True
while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        # 在此处绘制游戏内容
        if event.type == pygame.MOUSEBUTTONDOWN:  # 检测鼠标按下事件
            if event.button == 1:
                make = 1
        if event.type == pygame.MOUSEBUTTONUP:  # 检测鼠标抬起事件
            if event.button == 1:
                # 进行预测
                testing = np.expand_dims(test / 255.0, axis=0)  # 添加一个维度以适应模型输入
                predictions = model.predict(testing)
                predicted_label = np.argmax(predictions[0])  # 获取概率最大的类别索引
                # probabilities = predictions[0]
                probability = tf.nn.softmax(predictions[0]).numpy()  # 获取预测概率
                print("预测结果为：", predicted_label)
                print(probability)
                make = 0
        if make == 1:
            mouse_x, mouse_y = pygame.mouse.get_pos()
            mouse_x, mouse_y = amplitude_reduction(mouse_x, mouse_y, window_width, window_height)
            # print("鼠标x坐标:", mouse_x, "鼠标y坐标:", mouse_y)
            mouse_x = mouse_x // 20
            mouse_y = mouse_y // 20
            mouse_x, mouse_y = amplitude_reduction(mouse_x, mouse_y, 26, 26)
            pixel[mouse_x][mouse_y] = 1
            pixel[mouse_x + 1][mouse_y] = 1
            pixel[mouse_x][mouse_y + 1] = 1
            pixel[mouse_x + 1][mouse_y + 1] = 1
            test[mouse_y][mouse_x] = 255
        if event.type == pygame.MOUSEBUTTONDOWN:  # 检测鼠标按下事件
            if event.button == 3:
                pixel[:] = 0
                test[:] = 0
                predicted_label = 0
                probability[:] = 0
    window.fill((0, 0, 0))
    # 绘制
    for i in range(28):
        for j in range(28):
            if pixel[i][j] == 1:
                draw_stroke_rect(window, [i * 20, j * 20], rect_size, rect_outline, color_white, color_grey)
            else:
                draw_stroke_rect(window, [i * 20, j * 20], rect_size, rect_outline, color_black, color_grey)
    # 绘制文字
    text_surface = font.render('This number is:', True, color_white)  # 渲染文字
    window.blit(text_surface, (580, 10))  # 将文字渲染到指定位置

    text_num = font_num1.render(str(predicted_label), True, color_white)  # 渲染文字
    window.blit(text_num, (650, 55))  # 将文字渲染到指定位置

    for i in range(10):
        text_num = font_num2.render('Prob.' + str(i) + ': ' + str(round(probability[i] * 100, 2)) + '%', True, color_white)  # 渲染文字
        window.blit(text_num, (580, 120 + i * 40))  # 将文字渲染到指定位置

    pygame.display.update()  # 刷新显示窗口

# 退出Pygame
pygame.quit()
